import os
import cv2
def xywh2xyxy(x, imgsize):
    y = x.copy()
    WIDTH, HEIGHT = imgsize[0], imgsize[1]
    y[0] = x[0] - x[2] / 2  # top left x
    y[1] = x[1] - x[3] / 2  # top left y
    y[2] = x[0] + x[2] / 2  # bottom right x
    y[3] = x[1] + x[3] / 2  # bottom right y
    y[0] = y[0] * WIDTH + 1
    y[2] = y[2] * WIDTH + 1
    y[1] = y[1] * HEIGHT + 1
    y[3] = y[3] * HEIGHT + 1
    return y

def yolo2voc(content, root, file, categories, imgsize=(640, 640)):
    '''
    标注文件转换 txt转xml（YOLO to VOC）
    目录组织格式如下：
        rootpath-|
                 |-annotations-|
                 |             |-123.txt
                 |             |-124.txt
                 |             |-125.txt
                 |             |...
                 |
                 |-Cvoc--------|
                 |             |-123.xml
                 |             |-124.xml
                 |             |-125.xml
                 |             |...
                 |
                 |........
        其中, annotations是预先准备好的数据标签, Cvoc是自动创建的输出目录

    :param content:list[str], all lines in the txt file
    :param root: str, rootpath
    :param file: str, each .txt file
    :param categories: dict, {labelcode: 'labelname'}, eg: {0:'dog', 1:'bus'}
    :param imgsize: tuple, the original image size. imgsize=(WIDTH, HEIGHT)
    :return:
    '''
    savePath = os.path.join(root, "Cvoc")
    WIDTH, HEIGHT = imgsize[0], imgsize[1]
    if not os.path.exists(savePath):
        os.makedirs(savePath)

    with open(os.path.join(savePath, file[:-4]) + ".xml", 'w', encoding='utf-8') as xml:
        xml.write('<annotation>\n')
        xml.write('\t<folder>' + os.path.split(root)[1] + '</folder>\n')
        xml.write('\t<filename>' + file + '</filename>\n')
        xml.write('\t<path>' + os.path.join(root, file) + '</path>\n')
        xml.write('\t<size>\n')
        xml.write('\t\t<width>' + str(WIDTH) + '</width>\n')
        xml.write('\t\t<height>' + str(HEIGHT) + '</height>\n')
        xml.write('\t</size>\n')

        cnt = 0
        for line in content:
            line = line.split()
            cls, *bbox = line
            cls, bbox = int(cls), [float(i) for i in bbox]
            # bbox = xywh2xyxy(bbox, imgsize)

            xmin, ymin, xmax, ymax = bbox

            if cls==0:
                cls_name = categories[cls]
            else:
                continue

            xml.write('\t<object>\n')
            xml.write('\t\t<name>' + str(cls_name) + '</name>\n')
            xml.write('\t\t<pose>Unspecified</pose>\n')
            xml.write('\t\t<bndbox>\n')
            xml.write('\t\t\t<xmin>' + str(int(xmin)) + '</xmin>\n')
            xml.write('\t\t\t<ymin>' + str(int(ymin)) + '</ymin>\n')
            xml.write('\t\t\t<xmax>' + str(int(xmax)) + '</xmax>\n')
            xml.write('\t\t\t<ymax>' + str(int(ymax)) + '</ymax>\n')
            xml.write('\t\t</bndbox>\n')
            xml.write('\t</object>\n')
            cnt += 1

        if cnt == 0:
            print(file)
        xml.write('</annotation>')

def yolo2voc_by_file(content, root, file, categories, imgsize=(640, 640)):
    '''
    标注文件转换 txt转xml（YOLO to VOC）
    目录组织格式如下：
        rootpath-|
                 |-annotations-|
                 |             |-123.txt
                 |             |-124.txt
                 |             |-125.txt
                 |             |...
                 |
                 |-Cvoc--------|
                 |             |-123.xml
                 |             |-124.xml
                 |             |-125.xml
                 |             |...
                 |
                 |........
        其中, annotations是预先准备好的数据标签, Cvoc是自动创建的输出目录

    :param content:list[str], all lines in the txt file
    :param root: str, rootpath
    :param file: str, each .txt file
    :param categories: dict, {labelcode: 'labelname'}, eg: {0:'dog', 1:'bus'}
    :param imgsize: tuple, the original image size. imgsize=(WIDTH, HEIGHT)
    :return:
    '''
    savePath = os.path.join(root, "xml")
    WIDTH, HEIGHT = imgsize[0], imgsize[1]
    if not os.path.exists(savePath):
        os.makedirs(savePath)
    abpath, name = os.path.split(file)
    with open(os.path.join(savePath, name[:-4]) + ".xml", 'w', encoding='utf-8') as xml:
        xml.write('<annotation>\n')
        xml.write('\t<folder>' + os.path.split(root)[1] + '</folder>\n')
        image_path_jpg = os.path.join(root, 'JPEGImages', name[:-4] + '.jpg')
        image_path_png = os.path.join(root, 'JPEGImages', name[:-4] + '.png')
        image_path_bmp = os.path.join(root, 'JPEGImages', name[:-4] + '.bmp')
        image_path_BMP = os.path.join(root, 'JPEGImages', name[:-4] + '.BMP')
        if os.path.exists(image_path_jpg):
            xml.write('\t<filename>' + name[:-4] + ".jpg" + '</filename>\n')
            image_true_path = image_path_jpg
            image_tmp = cv2.imread(image_true_path)
            height, width, _ = image_tmp.shape
            if not WIDTH == width or not HEIGHT == height:
                print('\033[34m', "modify shape:", width, height, image_true_path)
                WIDTH = width
                HEIGHT = height
                imgsize = (WIDTH, HEIGHT)
        elif os.path.exists(image_path_png):
            xml.write('\t<filename>' + name[:-4] + ".png" + '</filename>\n')
            image_true_path = image_path_png
            image_tmp = cv2.imread(image_true_path)
            height, width, _ = image_tmp.shape
            if not WIDTH == width or not HEIGHT == height:
                print('\033[34m', "modify shape:", width, height, image_true_path)
                WIDTH = width
                HEIGHT = height
                imgsize = (WIDTH, HEIGHT)
        elif os.path.exists(image_path_bmp):
            xml.write('\t<filename>' + name[:-4] + ".bmp" + '</filename>\n')
            image_true_path = image_path_bmp
            image_tmp = cv2.imread(image_true_path)
            height, width, _ = image_tmp.shape
            if not WIDTH == width or not HEIGHT == height:
                print('\033[34m', "modify shape:", width, height, image_true_path)
                WIDTH = width
                HEIGHT = height
                imgsize = (WIDTH, HEIGHT)
        elif os.path.exists(image_path_BMP):
            xml.write('\t<filename>' + name[:-4] + ".BMP" + '</filename>\n')
            image_true_path = image_path_BMP
            image_tmp = cv2.imread(image_true_path)
            height, width, _ = image_tmp.shape
            if not WIDTH == width or not HEIGHT == height:
                print('\033[34m', "modify shape:", width, height, image_true_path)
                WIDTH = width
                HEIGHT = height
                imgsize = (WIDTH, HEIGHT)
        else:
            print('image {} read failed.'.format(os.path.join(savePath, name[:-4])))
            raise FileNotFoundError('Cant load image! Please check image path!')
        # xml.write('\t<filename>' + name[:-4] + "." + '</filename>\n')
        # xml.write('\t<path>' + os.path.join(root, name) + '</path>\n')
        xml.write('\t<path>' + image_true_path + '</path>\n')
        print('\034[34m', image_true_path)

        xml.write('\t<size>\n')
        xml.write('\t\t<width>' + str(WIDTH) + '</width>\n')
        xml.write('\t\t<height>' + str(HEIGHT) + '</height>\n')
        xml.write('\t\t<depth>' + str(1) + '</depth>\n')
        xml.write('\t</size>\n')

        cnt = 0
        for line in content:
            line = line.split()
            cls, *bbox = line
            cls, bbox = int(cls), [float(i) for i in bbox]
            bbox = xywh2xyxy(bbox, imgsize)

            xmin, ymin, xmax, ymax = bbox


            cls_name = categories[cls]


            xml.write('\t<object>\n')
            xml.write('\t\t<name>' + str(cls_name) + '</name>\n')
            xml.write('\t\t<pose>Unspecified</pose>\n')
            xml.write('\t\t<bndbox>\n')
            xml.write('\t\t\t<xmin>' + str(int(xmin)) + '</xmin>\n')
            xml.write('\t\t\t<ymin>' + str(int(ymin)) + '</ymin>\n')
            xml.write('\t\t\t<xmax>' + str(int(xmax)) + '</xmax>\n')
            xml.write('\t\t\t<ymax>' + str(int(ymax)) + '</ymax>\n')
            xml.write('\t\t</bndbox>\n')
            xml.write('\t</object>\n')
            cnt += 1

        if cnt == 0:
            print(file)
        xml.write('</annotation>')
# test example
if __name__ == "__main__":
    # output format by yolov5-detect.py is xywh and it has been normed

    # rootPath = "/data/ICE/DataSetCollection/record_b2/labeled/collection"
    # rootPath = "/data/ICE/DataSetCollection/indemind/collection/dataset"
    # rootPath = "/data/ICE/DataSetCollection/i18R_902_20210705/record_1/dataset/dataset_v1"
    # rootPath = "/data/ICE/DataSetCollection/i18R_902_20210705/record_2/dataset/dataset_v1"
    # rootPath = "/data/ICE/DataSetCollection/i18R_902_20210705/record_3/dataset/dataset_v1"
    # rootPath = "/data/ICE/DataSetCollection/batch2_batch3_tiny"
    # rootPath = "/data/ICE/DataSetCollection/batch4_06_04_extract"
    # rootPath = "/data/ICE/DataSetCollection/batch5_06_04_extract"
    # rootPath = "/data/ICE/DataSetCollection/batch6_20210610_extract/dataset"
    rootpath_list = [
        # '/data/ICE/DataSetCollection/record_b2/labeled/collection',
        #          '/data/ICE/DataSetCollection/indemind/collection/dataset',
        #          '/data/ICE/DataSetCollection/902_bug_20210826/collection',
        #          '/data/ICE/DataSetCollection/i18R_902_20210705/record_1/dataset/dataset_v1',
        #          '/data/ICE/DataSetCollection/i18R_902_20210705/record_2/dataset/dataset_v1',
        #          '/data/ICE/DataSetCollection/i18R_902_20210705/record_3/dataset/dataset_v1',
        #          '/data/ICE/DataSetCollection/batch2_batch3_tiny',
        #          '/data/ICE/DataSetCollection/batch4_06_04_extract',
        #          '/data/ICE/DataSetCollection/batch5_06_04_extract',
        #          '/data/ICE/DataSetCollection/batch6_20210610_extract/dataset',
        #          '/data/ICE/DataSetCollection/batch7_06_21_valid_extract/dataset',
        #          '/data/ICE/DataSetCollection/escalator_cover/collection',
        # '/data/ICE/DataSetCollection/i18Rperson_model_dummy/dataset_v1',
        # '/data/ICE/DataSetCollection/openData/all_usc_data',
        '/data/ICE/coco/indemind_train2017/VOC2007'
                 ]
    for rootPath_index in rootpath_list:
        # print(rootPath_index)
        labels_txt = os.path.join(rootPath_index, "all.txt")
        save_result_path = os.path.join(rootPath_index, "xml")

        files = []
        with open(labels_txt, "r") as file_list:
            lines = file_list.readlines()
            for line in lines:
                # print(line)
                text = line.strip()
                files.append(text)
        # print(files)
        # filePath = os.path.join(rootPath, "labels")
        # files = os.listdir(filePath)
        WIDTH = 640
        HEIGHT = 400
        categories = {0: 'person', 1: 'escalator', 2: 'escalator_handrails', 3: 'person_dummy', 4: 'escalator_model',
                      5: 'escalator_handrails_model'}

        for file in files:
            if '.txt' not in file:
                continue
            with open(file, "r") as f:
                abpath, name = os.path.split(file)
                root_abpath, _ = os.path.split(abpath)

                lines = f.readlines()
                yolo2voc_by_file(lines, root_abpath, file, categories, (WIDTH, HEIGHT))

